package git.soulbgm;

import cn.hutool.core.date.DateUtil;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.thread.ThreadUtil;
import cn.hutool.core.util.NumberUtil;
import cn.hutool.core.util.RandomUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.log.Log;
import cn.hutool.log.dialect.console.ConsoleLog;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.math.BigDecimal;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.nio.file.StandardOpenOption;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.Date;
import java.util.List;
import java.util.Random;
import java.util.concurrent.CountDownLatch;
import java.util.stream.Collectors;

import static cn.hutool.core.date.DatePattern.NORM_DATETIME_FORMATTER;

/**
 * 程序入口
 *
 * @author SoulBGM
 * @date 2024-08-30
 */
public class Application {

    private static final Log LOG = new ConsoleLog(Application.class);
    private static final Random RANDOM = new Random();

    public static final String LOCAL_OS;
    public static final String OS_WIN = "Win";

    static {
        LOCAL_OS = System.getProperty("os.name");
    }

    public static void main(String[] args) throws Exception {
        if (args.length != 3) {
            LOG.error("请添加执行参数, 参数1 总数 | 参数2 每个线程执行数量 | 参数3 数据生成根目录");
            System.exit(-1);
        }
        Path basePath = Paths.get(args[2]);
        if (Files.exists(basePath)) {
            FileUtil.del(basePath);
        }
        Files.createDirectory(basePath);

        long startTimestamp = DateUtil.parse("2023-01-01 00:00:00.000", DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSS")).getTime();

        int total = Integer.parseInt(args[0]);
        // I/O密集型任务，可以考虑核心数的两倍作为初始尝试
        int count = Runtime.getRuntime().availableProcessors();
        long fileLine = Long.parseLong(args[1]);
        List<Long> split = split(total, fileLine * count);

        long sumLine = 0;
        long startTime = System.currentTimeMillis();
        LOG.info("开始插入操作");
        for (int i = 0; i < split.size(); i++) {
            Long groupLine = split.get(i);
            List<Long> groupSplit = split(groupLine, fileLine);
            CountDownLatch latch = new CountDownLatch(groupSplit.size());
            List<String> nameList = new ArrayList<>(count);
            for (int j = 0; j < groupSplit.size(); j++) {
                Long batchSize = groupSplit.get(j);
                String name = "t_data_" + i + "_" + j + ".sql";
                ThreadUtil.execAsync(new WriteFileThread(Paths.get(basePath.toString(), name), batchSize, startTimestamp, latch));
                startTimestamp += batchSize;
                sumLine += batchSize;
                nameList.add(name);
            }
            latch.await();
            int seq = i + 1;
            LOG.info("第{}组写入完成, 当前写入行数{}", seq, sumLine);
            //execStr(StrUtil.format("cockroach userfile upload -r {} userfile:/// --certs-dir=/data/crdb/certs", basePath));
            //LOG.info("上传第{}组数据", seq);
            //String command = StrUtil.format("source /etc/profile;psql.py -t T_DATA_COPY no1:2181 {}/*.csv", basePath.toString());
            //cockroach userfile upload -r $PWD/t_data_copy userfile:/// --certs-dir=/data/crdb/certs
            //userfile://defaultdb.public.userfiles_root/csv/t_data_2_8.csv
            //cockroach userfile delete /csv  --certs-dir=/data/crdb/certs
            //cockroach sql --certs-dir=/data/crdb/certs -d tedb -e "IMPORT INTO t_data_hrj (id,name,pass,data,create_time,receive_time) CSV DATA('userfile://defaultdb.public.userfiles_root/dd.csv')"
            //String urls = nameList.stream().map(o -> "'userfile://defaultdb.public.userfiles_root/csv/" + o + "'").collect(Collectors.joining(","));
            //String command = StrUtil.format("source /etc/profile;cockroach sql --certs-dir=/data/crdb/certs -d tedb -e \"IMPORT INTO t_data_hrj (id,name,pass,data,create_time,receive_time) CSV DATA(" + urls + ")\"", basePath.toString());
            String command = StrUtil.format("for f in $(ls -l {}|grep -E .sql$|awk '{print $9}'); do cockroach sql --certs-dir=/data/crdb/certs -d souldb -f {}/$f;done", basePath.toString(), basePath.toString());
            execStr(command);
            LOG.info("导入第{}组数据", seq);
            //execStr("cockroach userfile delete /csv --certs-dir=/data/crdb/certs");
            //LOG.info("删除远程导入第{}组数据", seq);
            FileUtil.del(basePath);
            Files.createDirectory(basePath);
            LOG.info("删除第{}组数据文件", seq);
        }
        long endTime = System.currentTimeMillis();
        long diff = endTime - startTime;
        LOG.info("插入操作结束, 耗时 {} ms, 每秒平均插入 {} 条", diff, NumberUtil.div(total, NumberUtil.div(diff, 1000), 4));
    }

    /**
     * 将一个数以固定数进行分隔，不足则将剩余填充
     *
     * @param total 全部
     * @param num   数字
     * @return {@link List}<{@link Long}>
     */
    public static List<Long> split(long total, long num) {
        BigDecimal t = BigDecimal.valueOf(total);
        List<Long> splits = new ArrayList<>();
        long count = total / num;
        for (long i = 0; i < count; i++) {
            splits.add(num);
            t = t.subtract(BigDecimal.valueOf(num));
        }
        if (t.longValue() != 0) {
            splits.add(t.longValue());
        }
        return splits;
    }

    /**
     * 执行命令 如果执行错误也会将错误信息输出
     *
     * @param command 命令
     * @return 执行命令后返回的内容
     */
    public static String execStr(String command) {
        Process p;
        try {
            if (LOCAL_OS.contains(OS_WIN)) {
                p = new ProcessBuilder("cmd", "/c", command).redirectErrorStream(true).start();
            } else {
                p = new ProcessBuilder("/bin/bash", "-c", command).redirectErrorStream(true).start();
            }
        } catch (IOException e) {
            return null;
        }

        try (
                InputStream inputStream = p.getInputStream()
        ) {
            return getStreamOutput(inputStream);
        } catch (IOException e) {
            return null;
        }
    }

    /**
     * 得到流输出
     *
     * @param is 输入流
     * @return {@link String}
     * @throws IOException IO异常
     */
    private static String getStreamOutput(InputStream is) throws IOException {
        return getStreamOutput(is, LOCAL_OS.contains(OS_WIN) ? "GBK" : "UTF-8");
    }

    /**
     * 得到流输出
     *
     * @param is          输入流
     * @param charsetName 字符集名称
     * @return {@link String}
     * @throws IOException IO异常
     */
    private static String getStreamOutput(InputStream is, String charsetName) throws IOException {
        try (
                InputStreamReader in = new InputStreamReader(is, charsetName);
                BufferedReader br = new BufferedReader(in)
        ) {
            String line, output;
            StringBuilder sbLines = new StringBuilder();
            while ((line = br.readLine()) != null) {
                sbLines.append(line).append("\n");
            }
            output = sbLines.toString();
            return output;
        }
    }

    public static class WriteFileThread implements Runnable {

        private final Path path;
        private final long rows;
        private final long startTime;
        private final CountDownLatch latch;
        private static final String DATA = "11111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111";

        public WriteFileThread(Path path, long rows, long startTime, CountDownLatch latch) {
            this.path = path;
            this.rows = rows;
            this.startTime = startTime;
            this.latch = latch;
        }

        @Override
        public void run() {
            try {
                Files.createFile(path);
                StringBuilder sb = new StringBuilder();
                int length = String.valueOf(rows).length();
                String format = "%0" + length + "d";
                sb.append("INSERT INTO t_data_2 (id,name,pass,data,create_time,receive_time) VALUES ");
                for (long i = 1; i <= rows; i++) {
                    long receiveTime = startTime + i;
                    String line = //StrUtil.format("{},{},{},{},{},{}\n",
                            StrUtil.format("('{}','{}','{}','{}','{}','{}'),",
                            receiveTime + "" + String.format(format, i), getRandomChinese(RANDOM.nextInt(3) + 2),
                            getRandomString(RANDOM.nextInt(16) + 3), DATA,
                            //receiveTime + RANDOM.nextInt(100), receiveTime);
                            DateUtil.format(new Date(receiveTime + RANDOM.nextInt(100)), NORM_DATETIME_FORMATTER), DateUtil.format(new Date(receiveTime), NORM_DATETIME_FORMATTER));
                    sb.append(line);
                    if (i % 128 == 0) {
                        sb.deleteCharAt(sb.length() - 1);
                        sb.append(";\n");
                        Files.write(path, sb.toString().getBytes(), StandardOpenOption.APPEND);
                        sb.delete(0, sb.length());
                        sb.append("INSERT INTO t_data_2 (id,name,pass,data,create_time,receive_time) VALUES ");
                    }
                }
                sb.deleteCharAt(sb.length() - 1);
                sb.append(";\n");
                Files.write(path, sb.toString().getBytes(), StandardOpenOption.APPEND);
            } catch (IOException e) {
                e.printStackTrace();
            } finally {
                latch.countDown();
            }
        }
    }

    private static String getRandomString(int total) {
        return RandomUtil.randomString("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789", total);
    }

    private static String getRandomChinese(int total) {
        StringBuilder chineseStr = new StringBuilder();
        for (int i = 0; i < total; i++) {
            chineseStr.append(RandomUtil.randomChinese());
        }
        return chineseStr.toString();
    }

}
